import asyncio
import websockets
import threading
import subprocess
import os
from config import WEBSOCKET_HOST, WEBSOCKET_PORT


class WebSocketServer:
    def __init__(self):
        self.connected_clients = set()
        self.client_lock = threading.Lock()  # 新增：保护客户端集合的锁
        self.server = None
        self.loop = None  # 保存事件循环引用

    async def handler(self, websocket):
        """处理前端WebSocket连接（加锁保护集合操作）"""
        print("🟢 Web Agent 前端已连接")
        with self.client_lock:  # 加锁添加客户端
            self.connected_clients.add(websocket)
        try:
            async for message in websocket:
                # 可扩展：接收前端反馈
                pass
        except websockets.exceptions.ConnectionClosed:
            print("🔴 Web Agent 前端断开")
        finally:
            with self.client_lock:  # 加锁删除客户端
                self.connected_clients.discard(websocket)

    def get_server_url(self):
        return f"ws://{WEBSOCKET_HOST}:{WEBSOCKET_PORT}"

    def kill_port_if_occupied(self):
        """检查端口是否被占用，如果被占用则杀掉占用进程"""
        try:
            # 使用netstat查找占用指定端口的进程
            result = subprocess.run([
                'netstat', '-ano', '|', 'findstr', f':{WEBSOCKET_PORT}'
            ], capture_output=True, text=True, shell=True)
            
            lines = result.stdout.strip().split('\n')
            for line in lines:
                if 'LISTENING' in line:
                    parts = line.split()
                    if len(parts) >= 5 and parts[-1].isdigit():
                        pid = parts[-1]
                        print(f"⚠️ 检测到端口 {WEBSOCKET_PORT} 被PID {pid} 占用，正在终止...")
                        subprocess.run(['taskkill', '/PID', pid, '/F'], capture_output=True)
                        print(f"✅ 成功终止PID {pid}")
        except Exception as e:
            print(f"❌ 检查或终止端口占用进程时出错: {e}")

    def start(self):
        # 检查端口是否被占用，如果被占用则杀掉占用进程
        self.kill_port_if_occupied()
        
        # 在单独的线程中启动WebSocket服务器
        self.server_thread = threading.Thread(target=self.run_server, daemon=True)
        self.server_thread.start()
        print(f"🌐 WebSocket 服务器启动在 {self.get_server_url()}")

    def run_server(self):
        self.loop = asyncio.new_event_loop()  # 保存循环到实例变量
        asyncio.set_event_loop(self.loop)

        # 定义异步主函数，确保服务器在循环中启动
        async def main():
            # 使用websockets 15.x的新语法
            async with websockets.serve(
                self.handler,
                WEBSOCKET_HOST,
                WEBSOCKET_PORT
            ) as server:
                self.server = server
                await asyncio.Future()

        try:
            self.loop.run_until_complete(main())
        finally:
            self.loop.close()

    def get_connected_clients(self):
        with self.client_lock:  # 加锁读取集合
            return self.connected_clients.copy()  # 返回副本，避免外部修改

    def send_message(self, message):
        """向前端发送消息（加锁保护集合读取）"""
        with self.client_lock:  # 加锁检查并复制客户端
            if not self.connected_clients:
                print("⚠️ 无前端连接，无法发送消息")
                return
            clients = self.connected_clients.copy()  # 复制当前客户端列表

        async def _send_to_all():
            for client in clients:  # 迭代副本，避免原集合修改导致的问题
                try:
                    await client.send(message)
                    print(f"📤 已发送消息: {message}")
                except Exception as e:
                    print(f"❌ 发送失败: {e}")
                    with self.client_lock:  # 加锁删除无效客户端
                        self.connected_clients.discard(client)

        # 复用服务器的循环事件循环发送消息
        if self.loop and self.loop.is_running():
            asyncio.run_coroutine_threadsafe(_send_to_all(), self.loop)
        else:
            print("⚠️ 无法发送消息，服务器未启动")